import numpy as np

'''
自适应矩估计
'''

# 定义损失函数
def loss_function(y_pred, y_true):
    return np.mean((y_pred - y_true)**2)

# 定义模型
class Model:
    def __init__(self):
        self.w = np.random.randn()
        self.b = np.random.randn()
        self.m_w = 0
        self.m_b = 0
        self.v_w = 0
        self.v_b = 0
        self.beta1 = 0.9
        self.beta2 = 0.999
        self.epsilon = 1e-8
        self.t = 0

    def predict(self, x):
        return self.w*x + self.b

    def gradient(self, x, y_true):
        y_pred = self.predict(x)
        w_grad = 2*np.mean((y_pred - y_true)*x)
        b_grad = 2*np.mean(y_pred - y_true)
        return w_grad, b_grad

    # 使用Adam算法更新模型参数
    def update_parameters(self, w_grad, b_grad, learning_rate):
        self.t += 1
        self.m_w = self.beta1*self.m_w + (1 - self.beta1)*w_grad
        self.m_b = self.beta1*self.m_b + (1 - self.beta1)*b_grad
        self.v_w = self.beta2*self.v_w + (1 - self.beta2)*w_grad**2
        self.v_b = self.beta2*self.v_b + (1 - self.beta2)*b_grad**2
        m_w_hat = self.m_w/(1 - self.beta1**self.t)
        m_b_hat = self.m_b/(1 - self.beta1**self.t)
        v_w_hat = self.v_w/(1 - self.beta2**self.t)
        v_b_hat = self.v_b/(1 - self.beta2**self.t)
        self.w -= learning_rate*m_w_hat/(np.sqrt(v_w_hat) + self.epsilon)
        self.b -= learning_rate*m_b_hat/(np.sqrt(v_b_hat) + self.epsilon)

# 创建随机训练数据
x_train = np.linspace(0, 1, num=100)
y_train = 2*x_train + 1 + np.random.randn(100)*0.2

# 创建模型实例并使用Adam算法进行优化
model = Model()
learning_rate = 0.1
epochs = 50
for epoch in range(epochs):
    for i in range(len(x_train)):
        x = x_train[i]
        y_true = y_train[i]
        w_grad, b_grad = model.gradient(x, y_true)
        model.update_parameters(w_grad, b_grad, learning_rate)
    y_pred = model.predict(x_train)
    loss = loss_function(y_pred, y_train)
    print(f"Epoch: {epoch+1} | Loss: {loss:.4f}")